#pragma once

#include "../common/Common.h"
#include "../common/Base.h"
#include "../common/Message.h"

namespace rpc
{

    namespace server
    {

        class ProviderManager
        {
        public:
            using ptr = std::shared_ptr<ProviderManager>;

            struct Provider
            {
                using ptr = std::shared_ptr<Provider>;
                mutex mtx;
                BaseConnection::ptr conn;
                vector<string> methods;
                Address host;

                Provider() {}
                Provider(const BaseConnection::ptr &c, const Address &h) : conn(c), host(h) {}

                void appendMethod(const string &method)
                {
                    unique_lock<mutex> lock(mtx);
                    methods.emplace_back(method);
                }
            };
            // 当一个新的服务提供者进行服务注册的时候调用
            void addProvider(const BaseConnection::ptr &conn, const Address &host, const string &method)
            {
                Provider::ptr provider;

                {
                    // 找到提供者
                    unique_lock<mutex> lock(_mtx);
                    auto providerIt = _conns.find(conn);
                    if (providerIt == _conns.end())
                    {
                        provider = std::make_shared<Provider>(conn, host);
                        _conns.emplace(std::make_pair(conn, provider));
                    }
                    else
                    {
                        provider = providerIt->second;
                    }

                    // 如果unordered_map中不存在这个key， 那么会默认创建一个，其中value为默认值
                    auto &methods = _providers[method];
                    methods.emplace(provider);
                }

                // 向提供者添加信息
                provider->appendMethod(method);
            }

            // 当一个服务提供者断开连接的时候，删除它的关联信息
            void delProvider(const BaseConnection::ptr &conn)
            {
                // 找到方法提供者
                Provider::ptr provider;
                unique_lock<mutex> lock(_mtx);
                auto providerIt = _conns.find(conn);
                if (providerIt == _conns.end())
                {
                    // 不是方法提供者，直接返回
                    return;
                }

                provider = providerIt->second;
                for (auto &method : provider->methods)
                {
                    _providers[method].erase(provider);
                }

                _conns.erase(conn);
            }

            // 当一个服务提供者断开连接的时候，获取他的信息--用于进行服务下线通知
            Provider::ptr getProvider(const BaseConnection::ptr &conn)
            {
                unique_lock<mutex> lock(_mtx);
                auto it = _conns.find(conn);
                if (it == _conns.end())
                {
                    return Provider::ptr();
                }

                return it->second;
            }

            // 获取当前方法的所有提供者
            vector<Address> methodHosts(const string &method)
            {
                unique_lock<mutex> lock(_mtx);
                auto it = _providers.find(method);
                if (it == _providers.end())
                {
                    return vector<Address>();
                }

                vector<Address> ret;
                ret.reserve(it->second.size());
                for (auto &addr : it->second)
                {
                    ret.push_back(addr->host);
                }

                return ret;
            }

        private:
            mutex _mtx;
            unordered_map<string, set<Provider::ptr>> _providers;
            unordered_map<BaseConnection::ptr, Provider::ptr> _conns;
        };

        class DiscoverManager
        {
        public:
            using ptr = std::shared_ptr<DiscoverManager>;

            struct Discover
            {
                using ptr = std::shared_ptr<Discover>;
                mutex mtx;
                BaseConnection::ptr conn; // 发现者关联的客户端
                vector<string> methods;   // 发现的服务名称

                Discover() {}
                Discover(const BaseConnection::ptr &c) : conn(c) {}

                void appendMethod(const string &method)
                {
                    unique_lock<mutex> lock(mtx);
                    methods.emplace_back(method);
                }
            };

        public:
            // 当每次客户端进行服务发现的时候新增发现者，新增服务名称
            Discover::ptr addDiscover(const BaseConnection::ptr &conn, const string &method)
            {
                Discover::ptr discover;
                {
                    unique_lock<mutex> lock(_mtx);
                    auto connIt = _conns.find(conn);
                    if (connIt == _conns.end())
                    {
                        discover = std::make_shared<Discover>();
                        _conns.emplace(conn, discover);
                    }
                    else
                    {
                        discover = connIt->second;
                    }

                    auto &discovers = _discovers[method];
                    discovers.emplace(discover);
                }

                discover->appendMethod(method);
                return discover;
            }

            // 发现者客户端断开连接时，找到发现者信息，删除关联数据
            void delDiscover(const BaseConnection::ptr &conn)
            {
                unique_lock<mutex> lock(_mtx);
                auto connIt = _conns.find(conn);
                if (connIt == _conns.end())
                {
                    return;
                }

                for (auto &method : connIt->second->methods)
                {
                    _discovers[method].erase(connIt->second);
                }

                _conns.erase(conn);
            }

            // 当有一个新的服务提供者上线，则进行上线通知
            void onlineNotify(const string &method, const Address &host)
            {
                return notify(method, host, ServiceOpType::SERVICE_ONLINE);
            }

            // 当有一个服务提供者断开连接，则进行下线通知
            void offlineNotify(const std::string &method, const Address &host)
            {
                return notify(method, host, ServiceOpType::SERVICE_OFFLINE);
            }

        private:
            void notify(const string &method, const Address &host, ServiceOpType type)
            {

                auto it = _discovers.find(method);
                if (it == _discovers.end())
                {
                    return;
                }

                ServiceRequest::ptr req = MessageFactory::create<ServiceRequest>();
                req->setHost(host);
                req->setMtype(MType::REQ_SERVICE);
                req->setOpType(type);
                req->setRid(uuid());
                req->setMethod(method);

                for (auto &discover : it->second)
                {
                    discover->conn->send(req);
                }
            }

        private:
            mutex _mtx;
            unordered_map<string, set<Discover::ptr>> _discovers;
            unordered_map<BaseConnection::ptr, Discover::ptr> _conns;
        };

        class PDManager
        {
        public:
            using ptr = std::shared_ptr<PDManager>;

            PDManager()
                : _discover(std::make_shared<DiscoverManager>()), _provider(std::make_shared<ProviderManager>()) {}

            // 服务操作请求：服务注册/服务发现/
            void onServiceRequest(const BaseConnection::ptr &conn, ServiceRequest::ptr &req)
            {
                ServiceOpType type = req->opType();
                if (type == ServiceOpType::SERVICE_REGISTRY)
                {
                    ILOG("%s:%d 注册 %s 服务", req->host().first.c_str(), req->host().second, req->method().c_str());
                    _provider->addProvider(conn, req->host(), req->method());
                    _discover->onlineNotify(req->method(), req->host());
                    return registerResponse(conn, req);
                }
                else if (type == ServiceOpType::SERVICE_DISCOVER)
                {
                    ILOG("客户端要发现 %s 服务", req->method().c_str());
                    _discover->addDiscover(conn, req->method());
                    return discoverResponse(conn, req);
                }
                else
                {
                    ELOG("收到服务操作请求，但操作请求错误！");
                    return errorResponse(conn, req);
                }
            }

            // 连接断开时
            void onConnShutdown(const BaseConnection::ptr &conn)
            {
                ProviderManager::Provider::ptr provider = _provider->getProvider(conn);

                if (provider)
                {
                    ILOG("%s:%d 服务下线", provider->host.first.c_str(), provider->host.second);
                    // 获取当前提供者的所有提供方法
                    for (auto &method : provider->methods)
                    {
                        _discover->offlineNotify(method, provider->host);
                    }

                    // 删除当前订阅者
                    _provider->delProvider(conn);
                }

                // 从提供集合中删除
                _discover->delDiscover(conn);
            }

        private:
            void registerResponse(const BaseConnection::ptr &conn, ServiceRequest::ptr &req)
            {
                ServiceResponse::ptr resp = MessageFactory::create<ServiceResponse>();

                resp->setRid(req->rid());
                resp->setMtype(MType::RSP_SERVICE);
                resp->setOpType(ServiceOpType::SERVICE_REGISTRY);
                resp->setRcode(RCode::RCODE_OK);

                conn->send(resp);
            }

            void discoverResponse(const BaseConnection::ptr &conn, ServiceRequest::ptr &req)
            {
                ServiceResponse::ptr resp = MessageFactory::create<ServiceResponse>();
                resp->setRid(req->rid());
                resp->setMtype(MType::RSP_SERVICE);
                resp->setOpType(ServiceOpType::SERVICE_DISCOVER);

                vector<Address> hosts = _provider->methodHosts(req->method());
                if (hosts.empty())
                {
                    resp->setRcode(RCode::RCODE_NOT_FOUND_SERVICE);
                    return conn->send(resp);
                }

                resp->setHost(hosts);
                resp->setMethod(req->method());
                resp->setRcode(RCode::RCODE_OK);

                return conn->send(resp);
            }

            void errorResponse(const BaseConnection::ptr &conn, ServiceRequest::ptr &req)
            {
                ServiceResponse::ptr resp = MessageFactory::create<ServiceResponse>();
                resp->setRcode(RCode::RCODE_INVALID_OPTYPE);
                resp->setMtype(MType::RSP_SERVICE);
                resp->setOpType(ServiceOpType::SERVICE_UNKONW);
                resp->setRid(req->rid());

                conn->send(resp);
            }

        private:
            DiscoverManager::ptr _discover;
            ProviderManager::ptr _provider;
        };

    }; // server

}; // RPC